import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torchvision import models
from torch.utils.data import DataLoader, random_split, Subset
import os
import numpy as np
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from torch.optim.lr_scheduler import ReduceLROnPlateau

# 设置随机数种子
seed = 2024
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# 检查CUDA是否可用，并设置device变量
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像大小以适应VGG16模型
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.RandomRotation(10),      # 随机旋转
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 设置数据集路径
data_dir = r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\xgm\lab1_all\data\real_data'

# 加载数据集
full_dataset = torchvision.datasets.ImageFolder(root=data_dir, transform=transform)

# 获取每个类别的索引
class_indices = {}
for idx, (_, label) in enumerate(full_dataset):
    if label not in class_indices:
        class_indices[label] = []
    class_indices[label].append(idx)

# 划分训练集和测试集
train_indices = []
test_indices = []

for label, indices in class_indices.items():
    train_idx, test_idx = train_test_split(indices, test_size=0.2, random_state=seed)
    train_indices.extend(train_idx)
    test_indices.extend(test_idx)

train_dataset = Subset(full_dataset, train_indices)
test_dataset = Subset(full_dataset, test_indices)

# 创建数据加载器
train_loader = DataLoader(dataset=train_dataset, batch_size=4, shuffle=True, pin_memory=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=4, shuffle=False, pin_memory=True)

# 加载预训练的VGG16模型
vgg16 = models.vgg16(pretrained=True).to(device)

# 冻结卷积层
for param in vgg16.features.parameters():
    param.requires_grad = False

# 替换最后的全连接层以适应7分类任务
num_features = vgg16.classifier[6].in_features
vgg16.classifier[6] = nn.Linear(num_features, 7).to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(vgg16.classifier.parameters(), lr=0.001)
scheduler = ReduceLROnPlateau(optimizer, 'min', patience=5, factor=0.5, verbose=True)

# 初始化用于跟踪最佳模型的变量
best_accuracy = 0
loss_list = []
lr_list = []

# 训练和测试模型
num_epochs = 100
for epoch in range(num_epochs):
    vgg16.train()
    running_loss = 0.0
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)

        outputs = vgg16(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        if (i + 1) % 100 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}')

    epoch_loss = running_loss / len(train_loader)
    current_lr = optimizer.param_groups[0]['lr']
    loss_list.append(epoch_loss)
    lr_list.append(current_lr)

    print(f'Epoch [{epoch+1}/{num_epochs}] Loss: {epoch_loss:.4f}, Learning Rate: {current_lr:.6f}')

    # 调整学习率
    scheduler.step(epoch_loss)

    # 测试模型
    vgg16.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = vgg16(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        accuracy = 100 * correct / total
        print(f'Epoch [{epoch+1}/{num_epochs}] Test Accuracy: {accuracy:.2f}%')

        # 保存最好的模型权重
        if accuracy >= best_accuracy:
            best_accuracy = accuracy
            torch.save(vgg16.state_dict(), 'vgg16_best_model.pth')
            print("Saved Best Model Weights with Accuracy: {:.2f}%".format(accuracy))

# 加载最佳模型权重
vgg16.load_state_dict(torch.load('vgg16_best_model.pth'))

# 使用最佳模型对测试集进行测试
vgg16.eval()
correct = 0
total = 0
class_correct = [0 for _ in range(7)]
class_total = [0 for _ in range(7)]
misclassified_records = []

with torch.no_grad():
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = vgg16(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

        # 计算每个分类的正确和总数
        for i in range(len(labels)):
            label = labels[i]
            pred = predicted[i]
            if label == pred:
                class_correct[label] += 1
            class_total[label] += 1
            # 如果误判，记录图片路径、预测值和真实值
            if label != pred:
                img_path = full_dataset.samples[test_dataset.indices[total - len(labels) + i]][0]
                misclassified_records.append((img_path, pred.item(), label.item()))

accuracy = 100 * correct / total
print(f'Best Model Test Accuracy: {accuracy:.2f}%')

# 输出每个分类的准确率
for i in range(7):
    if class_total[i] > 0:
        class_accuracy = 100 * class_correct[i] / class_total[i]
        print(f'Class {i} Accuracy: {class_accuracy:.2f}%')

# 保存误判的图片路径、预测值和真实值到txt文件
with open('misclassified_records.txt', 'w') as f:
    for record in misclassified_records:
        f.write(f'{record[0]}\nPredicted: {record[1]}, Actual: {record[2]}\n')

print("Testing complete. Misclassified image paths and labels saved to 'misclassified_records.txt'.")

# 可视化Loss和学习率
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.plot(range(1, num_epochs + 1), loss_list, marker='o')
plt.title('Loss vs. Epochs')
plt.xlabel('Epochs')
plt.ylabel('Loss')

plt.subplot(1, 2, 2)
plt.plot(range(1, num_epochs + 1), lr_list, marker='o')
plt.title('Learning Rate vs. Epochs')
plt.xlabel('Epochs')
plt.ylabel('Learning Rate')

plt.tight_layout()
plt.savefig('training_metrics.png')
plt.show()
